"""
Code plots the traffic composition figure. Code is written in python 3.

by: Jacob Malone (j.malone@cablelabs.com)
date: August 25, 2021
"""


import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import rcParams
from matplotlib.patches import Rectangle

pd.options.display.float_format = lambda _: '%.2f' % _
pd.options.display.max_rows = 1000
pd.options.display.max_columns = 1000
pd.options.display.width = None

# Define working directory and load files from the "03_id_cord_cutters.py" files
data_dir = '/mnt/et01/cord_cutting_rep'

cord_cutters = pd.read_parquet(f'{data_dir}/04_cord_cutters.parquet')
other_subs = pd.read_parquet(f'{data_dir}/04_other_subs.parquet')
video_dates = pd.read_parquet(f'{data_dir}/04_video_dates.parquet')

allsubs = pd.concat([cord_cutters, other_subs], ignore_index=True)

cols_to_keep = ['date',
                'customer_key',
                'tot_gbdn',
                'tot_gbup',
                'vid_flag',
                'tier',
                'level_tunnel',
                'level_music',
                'level_gaming',
                'level_admin',
                'level_backup',
                'level_cdn',
                'level_other',
                'level_video',
                'level_sharing',
                'level_browsing',
                'level_streaming']
allsubs = allsubs.loc[:, cols_to_keep].copy()

vid_flag_max = (allsubs
                .groupby('customer_key')
                .vid_flag
                .max())
bund_accts = vid_flag_max.loc[vid_flag_max == 1].index
unbund_accts = vid_flag_max.loc[vid_flag_max == 0].index

bund_frame = allsubs.loc[allsubs.customer_key.isin(bund_accts), :]
unbund_frame = allsubs.loc[allsubs.customer_key.isin(unbund_accts), :]

# Focus on bundled customers first
level_vars = sorted([_ for _ in allsubs.columns if _.startswith('level_')])
bund_frame = (bund_frame
              .groupby('customer_key')[level_vars]
              .mean())
bund_frame = bund_frame.assign(tot_gb=bund_frame.loc[:, level_vars].sum(axis=1))
bund_frame = bund_frame.assign(bins=pd.qcut(bund_frame.tot_gb, 51, labels=False))
bund_bins_frame = (bund_frame
                   .groupby('bins')[level_vars + ['tot_gb']]
                   .sum())

prop_cols = []
for col in level_vars:
    pcol = 'prop_{0}'.format(col)
    prop_cols.append(pcol)
    bund_bins_frame[pcol] = bund_bins_frame[col] / bund_bins_frame.tot_gb

bund_bins_props = bund_bins_frame.loc[:, prop_cols].transpose()
bund_bins_dict = dict()
for col in bund_bins_props.columns:
    tmp_sum1 = (bund_bins_props.loc['prop_level_music', col] +
                bund_bins_props.loc['prop_level_streaming', col])
    tmp_sum2 = (bund_bins_props.loc['prop_level_admin', col] +
                bund_bins_props.loc['prop_level_backup', col] +
                bund_bins_props.loc['prop_level_cdn', col] +
                bund_bins_props.loc['prop_level_gaming', col] +
                bund_bins_props.loc['prop_level_sharing', col] +
                bund_bins_props.loc['prop_level_tunnel', col] +
                bund_bins_props.loc['prop_level_other', col])
    bund_bins_dict[col] = [bund_bins_props.loc['prop_level_video', col],
                            bund_bins_props.loc['prop_level_browsing', col],
                            tmp_sum1,
                            tmp_sum2]
    tmp_sum1, tmp_sum2 = None, None

bund_plt_frame = pd.DataFrame.from_dict(bund_bins_dict)
bund_plt_frame.index = ['video', 'browsing', 'music_streaming', 'other']


def plot_bund_area(plt_frame, filename):

    f, ax = plt.subplots()
    ax.xaxis.grid(color='black', linewidth=1)
    ax.yaxis.grid(color='black', linewidth=1)
    ax.set_axisbelow(False)

    X = plt_frame.columns
    Y = plt_frame.values

    areaplt = ax.stackplot(X, Y, baseline='zero', colors=['#595959',
                                                          '#FFFFFF',
                                                          '#B3B3B3',
                                                          '#000000'])

    ax.set_ylim([0, 1])
    ax.set_xlim([0, 50])
    ax.set_xlabel('Quantiles of Total Usage (Has Video Package)', size=20)
    ax.set_ylabel('Proportion of Total Usage', size=20)
    ax.set_xticks([0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50])
    ax.set_xticklabels([0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100])

    ax.tick_params(axis='x',
                   bottom='on',
                   top='off',
                   which='both',
                   labelsize=16)
    ax.tick_params(axis='y',
                   bottom='off',
                   top='off',
                   which='both',
                   labelsize=16)

    proxy_rects = [Rectangle((0, 0), 1, 1, fc=pc.get_facecolor()[0]) for pc in areaplt]
    legend = ax.legend(proxy_rects,
                       ('Video', 'Browsing', 'Music/Streaming', 'Other'),
                       fontsize=14,
                       frameon=True,
                       loc='lower right')
    legend.get_frame().set_linewidth(0.5)

    f.savefig('figs/{0}.png'.format(filename), bbox_inches='tight', dpi=300)

    plt.close(f)


# FIGURE A1, panel (a)
plot_bund_area(bund_plt_frame, 'cc_paper_bund_area_plot')

#####

# Focus on unbundled customers now
unbund_frame = (unbund_frame
                .groupby('customer_key')[level_vars]
                .mean())
unbund_frame = unbund_frame.assign(tot_gb=unbund_frame.loc[:, level_vars].sum(axis=1))
unbund_frame = unbund_frame.assign(bins=pd.qcut(unbund_frame.tot_gb, 51, labels=False))
unbund_bins_frame = (unbund_frame
                        .groupby('bins')[level_vars + ['tot_gb']]
                        .sum())

prop_cols = []
for col in level_vars:
    pcol = 'prop_{0}'.format(col)
    prop_cols.append(pcol)
    unbund_bins_frame[pcol] = unbund_bins_frame[col] / unbund_bins_frame.tot_gb

unbund_bins_props = unbund_bins_frame.loc[:, prop_cols].transpose()
unbund_bins_dict = dict()
for col in unbund_bins_props.columns:
    tmp_sum1 = (unbund_bins_props.loc['prop_level_music', col] +
                unbund_bins_props.loc['prop_level_streaming', col])
    tmp_sum2 = (unbund_bins_props.loc['prop_level_admin', col] +
                unbund_bins_props.loc['prop_level_backup', col] +
                unbund_bins_props.loc['prop_level_cdn', col] +
                unbund_bins_props.loc['prop_level_gaming', col] +
                unbund_bins_props.loc['prop_level_sharing', col] +
                unbund_bins_props.loc['prop_level_tunnel', col] +
                unbund_bins_props.loc['prop_level_other', col])
    unbund_bins_dict[col] = [unbund_bins_props.loc['prop_level_video', col],
                            unbund_bins_props.loc['prop_level_browsing', col],
                            tmp_sum1,
                            tmp_sum2]
    tmp_sum1, tmp_sum2 = None, None

unbund_plt_frame = pd.DataFrame.from_dict(unbund_bins_dict)
unbund_plt_frame.index = ['video', 'browsing', 'music_streaming', 'other']


def plot_unbund_area(plt_frame, filename):

    f, ax = plt.subplots()
    ax.xaxis.grid(color='black', linewidth=1)
    ax.yaxis.grid(color='black', linewidth=1)
    ax.set_axisbelow(False)

    X = plt_frame.columns
    Y = plt_frame.values

    areaplt = ax.stackplot(X, Y, baseline='zero', colors=['#595959',
                                                          '#FFFFFF',
                                                          '#B3B3B3',
                                                          '#000000'])

    ax.set_ylim([0, 1])
    ax.set_xlim([0, 50])
    ax.set_xlabel('Quantiles of Total Usage (No Video Package)', size=20)
    ax.set_ylabel('Proportion of Total Usage', size=20)
    ax.set_xticks([0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50])
    ax.set_xticklabels([0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100])

    ax.tick_params(axis='x',
                   bottom='on',
                   top='off',
                   which='both',
                   labelsize=16)
    ax.tick_params(axis='y',
                   bottom='off',
                   top='off',
                   which='both',
                   labelsize=16)

    proxy_rects = [Rectangle((0, 0), 1, 1, fc=pc.get_facecolor()[0]) for pc in areaplt]
    legend = ax.legend(proxy_rects,
                       ('Video', 'Browsing', 'Music/Streaming', 'Other'),
                       fontsize=14,
                       frameon=True,
                       loc='lower right')
    legend.get_frame().set_linewidth(0.5)

    f.savefig('figs/{0}.png'.format(filename), bbox_inches='tight', dpi=300)

    plt.close(f)


# FIGURE A1, panel (b)
plot_unbund_area(unbund_plt_frame, 'cc_paper_unbund_area_plot')